functor LinearScan (S : LINEAR_SCAN_STRUCTS) : LINEAR_SCAN =
struct
  open S
  
  structure Machine = Machine
 
  local open Machine
  in structure IdSet = IdSet
     structure Interval = Interval
     structure L = Location
     structure B = Binop
     structure R = Register
     structure O = Operand
     structure K = Block
     structure S = Structs
     structure J = Transfer
     structure M = Stm
     structure F = Function
     structure P = Program
     structure Intl = Intl
     structure Id = Id
     structure T = Types
  end
  
  val tag = false
  fun debugPrint s =
      if tag
      then print (String.concat [s, "\n"])
      else ()
      
  val doAgain = ref false
  
  val {get = liveIn, set = setLiveIn, destroy = destroyLiveIn, ...} :
      {get : Id.t -> IdSet.t, set : Id.t * IdSet.t -> unit, destroy : unit -> unit} =
      Property.destGetSet (Id.plist, Property.initConst IdSet.empty)
      
  val {get = liveOut, set = setLiveOut, destroy = destroyLiveOut, ...} :
      {get : Id.t -> IdSet.t, set : Id.t * IdSet.t -> unit, destroy : unit -> unit} =
      Property.destGetSet (Id.plist, Property.initConst IdSet.empty)
      
  val {get = liveInCounter, set = setLiveInCounter, destroy = destroyLiveInCounter, ...} :
      {get : Counter.t -> IdSet.t, set : Counter.t * IdSet.t -> unit, destroy : unit -> unit} =
      Property.destGetSet (Counter.plist, Property.initConst IdSet.empty)
      
  val {get = liveOutCounter, set = setLiveOutCounter, destroy = destroyLiveOutCounter, ...} :
      {get : Counter.t -> IdSet.t, set : Counter.t * IdSet.t -> unit, destroy : unit -> unit} =
      Property.destGetSet (Counter.plist, Property.initConst IdSet.empty)
      
  val {get = interval, set = setInterval, destroy = destroyInterval, ...} :
      {get : Id.t -> Interval.t, set : Id.t * Interval.t -> unit, destroy : unit -> unit} =
      Property.destGetSet (Id.plist, Property.initConst (Interval.bogus ()))
      
  val {get = allocation, set = setAllocation, destroy = destroyAllocation, ...} :
      {get : Id.t -> L.t , set : Id.t * L.t -> unit, destroy : unit -> unit} =
      Property.destGetSetOnce (Id.plist, Property.initConst (L.bogus ()))
      
      
  fun destroy () =
     (destroyLiveIn ()
     ; destroyLiveOut ()
     ; destroyLiveInCounter ()
     ; destroyLiveOutCounter ()
     ; destroyInterval ()
     ; destroyAllocation ())
  
  structure LivenessBlock :
  sig
    val main : P.t -> unit
  end =
  struct
  fun doit bs =
      case bs
       of [] => ()
        | x :: xs =>
          let val name = K.name x
              val succs = K.successors x
              val liveIns = List.map liveIn succs
              val out = List.foldl (fn (a, b) => IdSet.union (a, b)) IdSet.empty liveIns
              val liveOutNow = liveOut name
          in  if IdSet.equals (out, liveOutNow)
              then doit xs
              else let val _ = doAgain := true
                       val _ = setLiveOut (name, out)
                       val genX = K.gen x
                       val killX = K.kill x
                       val _ = setLiveIn (name, IdSet.union (genX, IdSet.difference (out, killX)))
                   in  ()
                   end
          end
          
  fun transBlocks bs =
      let val _ = doit bs
      in  if !doAgain
          then let val _ = doAgain := false
               in  transBlocks bs
               end
          else ()
      end
  
  fun initialize bs =
      List.app (fn x => 
                   let val gen = K.gen x
                       val name = K.name x
                   in  setLiveIn (name, gen);
                       setLiveOut (name, IdSet.empty)
                   end)
               bs
  fun transFunc f = 
      case f
       of F.T {name, exit, rettype, args, locals, stackSize, blocks, pre, post} =>
          let val _ = initialize blocks
          in  transBlocks blocks
          end
    
  fun checkBlock b =
      let val name = K.name b
          val liveIn = liveIn name
          val liveOut = liveOut name
          val strLiveIn = IdSet.toString liveIn
          val strLiveOut = IdSet.toString liveOut
          val _ = debugPrint (String.concat ["the live in set for ", Id.toString name, " is:"])
          val _ = debugPrint strLiveIn
          val _ = debugPrint (String.concat ["the live out set for ", Id.toString name, " is:"])
          val _ = debugPrint strLiveOut
      in  ()
      end
  fun checkBlocks bs = List.app checkBlock bs
  
  fun checkFunc f = 
      case f
       of F.T {name, exit, rettype, args, locals, stackSize, blocks, pre, post} =>
          checkBlocks blocks
    
  fun transProg p =
      case p
       of P.T {structs,
               gvars,
               funcs} =>
          let val _ = List.app transFunc funcs
              val _ = if tag
                      then List.app checkFunc funcs
                      else ()
          in  ()
          end
          
  val main = transProg
  end
  
  structure LivenessStm :
  sig
    val main : P.t -> unit
  end =
  struct  
  fun calculateStm (ss, liveOutForS, sentinal) =
      case ss
       of [] =>
          if IdSet.equals (liveOutForS, sentinal)
          then ()
          else raise Fail "compiler bug"
        | x :: xs =>
          let val _ = setLiveOutCounter (M.line x, liveOutForS)
              val liveInForS = IdSet.union (IdSet.fromList (M.gen x),
                                            IdSet.difference (liveOutForS,
                                                              IdSet.fromList (M.kill x)))
              val _ = setLiveInCounter (M.line x, liveInForS)
          in  calculateStm (xs, liveInForS, sentinal)
          end
  
  fun calculateTransfer (t, liveOutForB) =
      case t
       of J.Jump {line, label} =>
          let val _ = setLiveOutCounter (line, liveOutForB)
              val liveInForB = liveOutForB
              val _ = setLiveInCounter (line, liveInForB)
          in  liveInForB
          end
        | J.If {line, value, truee, falsee} =>
          let val _ = setLiveOutCounter (line, liveOutForB)
              val liveInForB = IdSet.union (liveOutForB, IdSet.fromList (J.gen t))
              val _ = setLiveInCounter (line, liveInForB)
          in  liveInForB         
          end
        | J.Return {line, value} =>
          let val _ = setLiveOutCounter (line, liveOutForB)
              val liveInForB = IdSet.union (liveOutForB, IdSet.fromList (J.gen t))
              val _ = setLiveInCounter (line, liveInForB)
          in  liveInForB          
          end
  
  fun transBlock b =
      case b
       of K.T {name, prop, stms, transfer, successors, ...} =>
          let val liveInForB = liveIn name
              val liveOutForB = liveOut name
              val new = calculateTransfer (transfer, liveOutForB)
              val _ = calculateStm (List.rev stms, new, liveInForB)
          in  ()
          end
  
  fun transFunc f = 
      case f
       of F.T {name, exit, rettype, args, locals, stackSize, blocks, pre, post} =>
          let val _ = List.app transBlock blocks
          in  ()
          end
         
  fun transProg p =
      case p
       of P.T {structs,
               gvars,
               funcs} =>
          let val _ = List.app transFunc funcs
              (*
              val _ = if tag
                      then checkFuncs funcs
                      else ()
                      *)
          in  ()
          end
          
  val main = transProg
  end
  
  structure CalculateInterval =
  struct
  fun doit (l, line) =
      case l
       of [] => ()
        | x :: xs =>
          let val old = interval x
              val new = Interval.shrink (old, line)
              val _ = setInterval (x, new)
          in  doit (xs, line)
          end
          
  fun transStm s =
      let val line = M.line s
          val lives = liveInCounter line
      in  doit (IdSet.toList lives, line)
      end
          
  fun transTransfer t =
      let val line = J.line t
          val lives = liveInCounter line
      in  doit (IdSet.toList lives, line)
      end
  
  fun transBlock b =
      case b
       of K.T {name, prop, stms, transfer, successors, ...} =>
          let val _ = transTransfer transfer
              val _ = List.app transStm stms
          in  ()
          end
      
  fun transFunc f = 
      case f
       of F.T {name, exit, rettype, args, locals, stackSize, blocks, pre, post} =>
          let val _ = List.app transBlock blocks
          in  ()
          end
    
  fun transProg p =
      case p
       of P.T {structs,
               gvars,
               funcs} =>
          let val _ = List.app transFunc funcs
              (*
              val _ = if tag
                      then List.app checkFunc funcs
                      else ()
                      *)
          in  ()
          end
          
  fun main p = transProg p
  end
  
  structure LinearScan :
  sig
    val main : P.t -> P.t
  end =
  struct
  structure AllIntervals :
  sig
    val all : {id : Id.t, loc : L.t, ty : T.t} list ref
    
    val init : {id : Id.t, loc : L.t, ty : T.t} list -> unit
    val isEmpty : unit -> bool
    val getOne : unit -> {id : Id.t, loc : L.t, ty : T.t}
  end =
  struct    
    val all : {id : Id.t, loc : L.t, ty : T.t} list ref = ref []
    
    fun init s = all := s
    fun isEmpty () =
        case !all
         of [] => true
          | _ => false
    fun getOne () =
        case !all
         of [] => raise Fail "compiler bug"
          | x :: xs =>
            let val _ = all := xs
            in  x
            end
  end 
  
  structure FreeRegs :
  sig
    val all : R.t list ref
    
    val addOne : R.t -> unit
    val init : unit -> unit
    val isEmpty : unit -> bool
    val getOne : unit -> R.t
  end =
  struct    
    val all : R.t list ref = ref []
    
    fun init () = all := R.allocatables ()
    fun isEmpty () =
        case !all
         of [] => true
          | _ => false
    fun getOne () =
        case !all
         of [] => raise Fail "compiler bug"
          | x :: xs =>
            let val _ = all := xs
            in  x
            end
    
    fun addOne r = all := r :: (!all)
  end
  
  structure ActiveIntervals :
  sig
    val all : {id : Id.t, loc : L.t, ty : T.t} list ref
    
    val addOne : {id : Id.t, loc : L.t, ty : T.t} -> unit
    val expire : Interval.t -> unit
    val init : unit -> unit
    val isEmpty : unit -> bool
    val getOne : unit -> {id : Id.t, loc : L.t, ty : T.t}
  end =
  struct    
    val all : {id : Id.t, loc : L.t, ty : T.t} list ref = ref []
    
    fun addOne (x as {id, loc, ty}) =
        let val a = !all
            fun doit l =
                case l
                 of [] => [x]
                  | (y as {id = idy, loc = locY, ty = tyy}) :: ys =>
                    let val (intervalX, intervalY) = (interval id, interval idy)
                    in  case Interval.tailCompare (intervalX, intervalY)
                         of LESS => x :: y :: ys
                          | EQUAL => x :: y :: ys
                          | GREATER => y :: (doit ys)
                    end
            val a = doit a
            val _ = all := a
        in  ()
        end
        
    fun init () = all := []
    fun isEmpty () =
        case !all
         of [] => true
          | _ => false
    fun getOne () =
        case !all
         of [] => raise Fail "compiler bug"
          | x :: xs =>
            let val _ = all := xs
            in  x
            end
            
    fun expire i =
        let val a = !all
            fun doit l =
                case l
                 of [] => []
                  | {id, loc, ty} :: xs =>
                    let val inter = interval id
                        val reg = allocation id
                        val reg = (case reg
                                   of L.InReg r => r
                                    | _ => raise Fail "compiler bug")
                    in  if Interval.expire (inter, i)
                        then let val _ = FreeRegs.addOne reg
                             in  doit xs
                             end
                        else l
                    end
            val a = doit a
            val _ = all := a 
        in  ()
        end
  end
  
  fun regAlloc locals =
      let val _ = debugPrint "regAlloc starting\n"
          val _ = AllIntervals.init locals
          val _ = FreeRegs.init ()
          val _ = ActiveIntervals.init ()
          fun doit () =
              if AllIntervals.isEmpty ()
              then ()
              else let val newInter as {id, loc, ty} = AllIntervals.getOne ()
                       val _ = ActiveIntervals.expire (interval id)
                       val _ = if FreeRegs.isEmpty ()
                               then let val _ = AllIntervals.getOne ()
                                    in  doit ()
                                    end
                               else let val newReg = FreeRegs.getOne ()
                                        val _ = debugPrint (String.concat ["do allocation:\n",
                                                                           Id.toString id,
                                                                           " ====> ",
                                                                           R.toString newReg])
                                        val _ = setAllocation (id, L.InReg newReg)
                                        val _ = ActiveIntervals.addOne newInter
                                    in  ()
                                    end
                   in  doit ()
                   end
          val _ = doit ()
      in  ()
      end
  
  fun sortLocals locals =
      let val locals = ListMergeSort.sort (fn ({id = id1, loc = l1, ty = ty1}, {id = id2, loc, ty}) =>
                                              let val (x, y) = (interval id1, interval id2)
                                              in  case Interval.headCompare (x, y)
                                                   of GREATER => true
                                                    | EQUAL => true
                                                    | LESS => false
                                              end)
                                          locals
      in  locals
      end
  
  structure Rewrite =
  struct
    val localOffSet = ref (~16)
    val argOffSet = ref 8
    
    fun init () = (localOffSet := ~16; argOffSet := 8)
    
    fun getLocal size =
        let val off = !localOffSet
            val _ = localOffSet := off - size            
        in  off
        end
        
    fun getArg () =
        let val off = !argOffSet
            val _ = argOffSet := off + 4
        in  off
        end
        
    fun rewriteLocals locals =
      List.map (fn {id, loc, ty} => 
                   let val l = allocation id
                       val l = case l 
                                of L.InReg r => L.InReg r
                                 | L.InStack _ => raise Fail "compiler bug"
                                 | L.Global => raise Fail "to do"
                                 | L.NoWhere =>
                                   let val size = 4
                                       val pos = getLocal size
                                   in  L.InStack (Intl.fromInt pos)
                                   end
                   in  {id = id, loc = l, ty = ty}
                   end)
              locals
              
    fun rewriteArgs args =
        List.map (fn {id, loc, ty} =>
                     {id = id,
                      loc = L.InStack (Intl.fromInt (getArg ())),
                      ty = ty})
                 args
  end
  
  
                    
  fun transFunc f = 
      case f
       of F.T {name, exit, rettype, args, locals, stackSize, blocks, pre, post} =>
          let val locals = sortLocals locals
              val _ = debugPrint (String.concat ["after sort locals in function: ", Id.toString name, "\n"])
              val _ = regAlloc locals
              val _ = Rewrite.init ()
              val locals = Rewrite.rewriteLocals locals
              val args = Rewrite.rewriteArgs args
              val new = F.T {name = name,
                             exit = exit,
                             rettype = rettype,
                             args = args,
                             locals = locals,
                             stackSize = Intl.abs (Intl.fromInt (!Rewrite.localOffSet + 16)),
                             blocks = blocks,
                             pre = pre,
                             post = post}
              val _ = destroyAllocation ()
          in  new
          end
          
  fun transProg p =
        case p
         of P.T {structs,
                 gvars,
                 funcs} =>
          let val funcs = List.map transFunc funcs
          in  P.T {structs = structs,
                   gvars = gvars,
                   funcs = funcs}
          end
          
    val main = transProg
  end
    
  fun trans p =
      let val _ = debugPrint "do liveness for blocks\n"
          val _ = LivenessBlock.main p
          val _ = debugPrint "do liveness for stms\n"
          val _ = LivenessStm.main p
          val _ = debugPrint "calculating intervals\n"
          val _ = CalculateInterval.main p
          val _ = debugPrint "linear scan starting\n"
          val p = LinearScan.main p
          val _ = destroy ()
      in  p
      end
end
